import matplotlib.pyplot as plt
import torch.nn.functional as F
import torch.nn as nn
from  data_process import *

class RNNModel(nn.Module):
    def __init__(self):
        super(RNNModel, self).__init__()
        self.embedding = nn.Embedding(encoder_number_embedding, encoder_embedding_dim,padding_idx=0)
        self.lstm = nn.LSTM(encoder_input_size, hidden_size=encoder_hidden_size,dropout=0.1, batch_first=True )
        self.fc1 = nn.Linear(in_features=10240,out_features=1024)
        self.fc2 = nn.Linear(in_features=1024,out_features=2)

    def forward(self, x):
        embeds = self.embedding(x)
        lstm_out, _ = self.lstm(embeds) #64,40,256
        # lstm_out2 = lstm_out[:, -1, :]
        lstm_out2 = lstm_out.reshape(batch_size, -1)
        out = self.fc1(lstm_out2)
        out = F.relu(out)
        out = self.fc2(out)
        out = F.log_softmax(out, dim=1)
        return out  #batch_size,2

model=RNNModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

train_loss=[]
train_precision=[]
for i in range(epochs):
    epoch_loss=0
    train_num = 0
    for inputs, labels,_ in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)
        model.zero_grad()
        output = model(inputs)
        loss = criterion(output, torch.squeeze(labels))
        epoch_loss+=loss.item()
        pre=torch.argmax(output, dim=1)
        train_num += (pre == torch.squeeze(labels)).sum().item()
        loss.backward()
        optimizer.step()
    train_loss.append(epoch_loss)
    train_precision.append(train_num / len(trainloader.dataset))
    print(f"({i + 1}/{epochs})训练准确率：{train_num / len(trainloader.dataset)}  ,训练损失值：{epoch_loss}")

print("-------------------测试------------------------")
model.eval()
test_nums=0
with torch.no_grad():
    for x,y,_ in testloader:
        x,y=x.to(device),y.to(device)
        pred = model(x)
        predicted = torch.argmax(pred, dim=-1)
        test_nums+=(predicted==torch.squeeze(y)).sum().item()
precision=test_nums/len(testloader.dataset)
print("test_precision:",precision)
torch.save(model, "rnn_model.pth")


plt.figure(figsize=(10,6))
plt.plot(range(1,epochs+1),train_loss)
plt.figure(figsize=(10,6))
plt.plot(range(1,epochs+1),train_precision)
plt.show()